import cv2 as cv
import numpy as np



def shape_detec(img, max_edge=5, arc_thresh=100):
    '''
    img processing, find contours, select, approx, drawrectangles and put labels
    '''
    bilateral = cv.bilateralFilter(img, d=15, sigmaColor=35, sigmaSpace=25)
    gray = cv.cvtColor(bilateral, cv.COLOR_BGR2GRAY)
    gray = cv.morphologyEx(gray, cv.MORPH_CLOSE, cv.getStructuringElement(cv.MORPH_RECT, (4,4)))
    thresh, thresholded = cv.threshold(gray, 220, 255, cv.THRESH_BINARY)

    canny = cv.Canny(thresholded, 125, 175)
    contours, hierarchies = cv.findContours(canny, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)
    contours_selected = list(filter(lambda e: cv.arcLength(e, True)>arc_thresh, contours)) 
    
    approx = []
    for cnt in contours_selected:
        eps = 0.02*cv.arcLength(cnt, True)
        appr = cv.approxPolyDP(cnt, eps, True)
        approx.append(appr) 
        
        top_point_x = 10000
        top_point_y = 10000
        bottom_point_x = 0
        bottom_point_y = 0
        for e in appr:
            top_point_x = min(top_point_x, e[0][0])
            top_point_y = min(top_point_y, e[0][1])
            bottom_point_x = max(bottom_point_x, e[0][0])
            bottom_point_y = max(bottom_point_y, e[0][1])
        cv.rectangle(img, (top_point_x, top_point_y), (bottom_point_x, bottom_point_y), (0, 0, 0), 2)
        if len(appr)==3:
            shape = 'Triangle'
        elif len(appr)==4:
            shape = 'Rectangle'
        elif len(appr)>max_edge:
            shape = 'Circle'
        cv.putText(img, shape, (top_point_x-3, top_point_y), cv.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 0), thickness = 1)

    # cv.drawContours(img, approx, -1, (0, 0, 255), 2, cv.LINE_AA)
    cv.imshow('selected', img)



def rescaleImg(img, scale=0.5):
    width = int(img.shape[1]*scale)
    height = int(img.shape[0]*scale)
    dimensions = (width, height)
    return cv.resize(img, dimensions, interpolation=cv.INTER_AREA)



img = cv.imread('shape.jpg')
img = rescaleImg(img, 1)
shape_detec(img)



cv.waitKey(0)